// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "sandboxed_api/tools/clang_generator/sandboxed_library_emitter.h"

#include <algorithm>
#include <cstddef>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "clang/AST/Decl.h"
#include "clang/AST/Mangle.h"
#include "clang/AST/Type.h"
#include "sandboxed_api/tools/clang_generator/emitter_base.h"
#include "sandboxed_api/tools/clang_generator/generator.h"

namespace sapi {
namespace {

constexpr absl::string_view kIncludePrefix = "";

// TODO(dvyukov): try to prepend a C++ namespace to C++ names
// instead of just adding a prefix. It will make it nicer
// for debuggers/profilers that demangle symbols.
constexpr absl::string_view kWrapperPrefix = "sapi_wrapper_";

// We pre-include some common headers always to avoid more complex logic
// to figure out when we actually need each of them.
// We have to disable whitespace/line_length linter checks because
// some identifiers are longer than the DEC VT52 terminal line length limit.
// We have to disable runtime/int linter checks because after typedef
// resolution we can inject types that are prohibited by linter,
// or if just the library used prohibited types.
constexpr absl::string_view kCommonHeader =
    R"(// Automatically generated by SAPI interface generator; DO NOT EDIT.

// NOLINTBEGIN(whitespace/line_length)
// NOLINTBEGIN(runtime/int)

#include <stdbool.h>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <cstdlib>
)";

constexpr absl::string_view kCommonFooter =
    R"(
// NOLINTEND(whitespace/line_length)
// NOLINTEND(runtime/int)
)";

constexpr absl::string_view kHeaderHeader =
    R"(// NOLINT(build/header_guard)
#pragma once

)";

constexpr absl::string_view kHostHeader = R"(
#include <memory>

#include "$0absl/log/check.h"
#include "$0sandboxed_api/vars.h"
#include "$0sandboxed_api/sandbox.h"

__attribute__((weak))
std::unique_ptr<sandbox2::Policy> $1SandboxModifyPolicy(
    sandbox2::PolicyBuilder* builder) {
  return builder->BuildOrDie();
}

struct $1SandboxImpl : public $1Sandbox {
  static $1SandboxImpl* Instance() {
    static $1SandboxImpl instance;
    return &instance;
  }

  $1SandboxImpl() { Check(Init()); }

  void Check(const absl::Status& status) {
    CHECK_OK(status) << "SAPI sandbox $1 failed";
  }

  std::unique_ptr<sandbox2::Policy> ModifyPolicy(
      sandbox2::PolicyBuilder* builder) override {
    return $1SandboxModifyPolicy(builder);
  }
};

)";

}  //  namespace

// Base class for different types of arguments.
// It provides an interface that allows to generate sandboxee/host wrappers
// with different number of actual arguments (e.g. a single argument
// in the original interface can be disassembled into/assembled from
// multiple arguments in the sandboxee wrapper interface)
// without knowing details of all possible argument types.
class SandboxedLibraryEmitter::Arg {
 public:
  Arg(absl::string_view name, absl::string_view type)
      : name_(name), type_(type) {}

  const std::string& EmitRetType() const { return type_; }
  std::string EmitHostParams() const {
    return absl::Substitute("$0 $1", type_, name_);
  }

  virtual std::vector<std::string> Includes() const { return {}; };
  virtual std::string EmitHostPreCall() const { return ""; }
  virtual std::string EmitHostPostCall() const { return ""; }
  virtual std::string EmitHostArgs() const = 0;
  virtual std::string EmitSandboxeeParams() const = 0;
  virtual std::string EmitSandboxeeArgs() const = 0;
  virtual std::string EmitSandboxeePreCall() const { return ""; }
  virtual std::string EmitSandboxeePostCall() const { return ""; }
  // TODO(dvyukov): Currently we pass return arguments as an additional argument
  // always to simplify the code. However, we could return scalar return values
  // directly since SAPI supports that, and that will be more efficient.
  virtual std::string EmitRetParams() const { LOG(FATAL) << "not implemented"; }
  virtual std::string EmitRetPreCall() const {
    LOG(FATAL) << "not implemented";
  }
  virtual std::string EmitRetArgs() const { LOG(FATAL) << "not implemented"; }
  virtual std::string EmitHostRet() const { LOG(FATAL) << "not implemented"; }
  virtual std::string EmitSandboxeeRet() const {
    LOG(FATAL) << "not implemented";
  }
  virtual ~Arg() = default;

 protected:
  const std::string name_;
  const std::string type_;
};

// Simple scalar arguments (ints).
// Nothing to see here, move along.
struct ScalarArg : SandboxedLibraryEmitter::Arg {
  using Arg::Arg;
  std::string EmitHostArgs() const override { return name_; }
  std::string EmitSandboxeeParams() const override { return EmitHostParams(); }
  std::string EmitSandboxeeArgs() const override { return name_; }
  std::string EmitRetParams() const override {
    return absl::Substitute("$0* $1", type_, name_);
  }
  std::string EmitRetPreCall() const override {
    return absl::Substitute("sapi::v::Reg<$0> $1_tmp;\n", type_, name_);
  }
  std::string EmitRetArgs() const override {
    return absl::Substitute("$0_tmp.PtrAfter()", name_);
  }
  std::string EmitHostRet() const override {
    return absl::Substitute("return $0_tmp.GetValue();\n", name_);
  }
  std::string EmitSandboxeeRet() const override {
    return absl::Substitute("*$0 = sapi_ret_val;\n", name_);
  }
};

// "const std::string&", these are always "input" to the library.
struct StringConstRefArg : SandboxedLibraryEmitter::Arg {
  using Arg::Arg;
  std::vector<std::string> Includes() const override { return {"<string>"}; }
  std::string EmitHostPreCall() const override {
    return absl::Substitute(
        "sapi::v::Array<const char> sapi_tmp_$0($0.data(), $0.size());\n",
        name_);
  }
  std::string EmitHostArgs() const override {
    return absl::Substitute("sapi_tmp_$0.PtrBefore(), $0.size()", name_);
  }
  std::string EmitSandboxeeParams() const override {
    return absl::Substitute("const char* $0_data, size_t $0_size", name_);
  }
  std::string EmitSandboxeeArgs() const override {
    return absl::Substitute("std::string($0_data, $0_size)", name_);
  }
};

// "std::string", same as "const std::string&", but also can be a return type.
struct StringArg : StringConstRefArg {
  using StringConstRefArg::StringConstRefArg;
  std::string EmitRetPreCall() const override {
    return "sapi::v::LenVal sapi_ret_tmp(0);\n";
  }
  std::string EmitRetArgs() const override { return "sapi_ret_tmp.PtrAfter()"; }
  virtual std::string EmitRetParams() const {
    return absl::Substitute("sapi::LenValStruct* $0", name_);
  }
  std::string EmitHostRet() const override {
    return "return "
           "std::string(reinterpret_cast<char*>(sapi_ret_tmp.GetData()), "
           "sapi_ret_tmp.GetDataSize());";
  }
  std::string EmitSandboxeeRet() const override {
    return absl::Substitute(
        "$0->data = strdup(sapi_ret_val.c_str());\n"
        "$0->size = sapi_ret_val.size();\n",
        name_);
  }
};

// "std::string&", input/output argument.
struct StringRefArg : SandboxedLibraryEmitter::Arg {
  using Arg::Arg;
  std::vector<std::string> Includes() const override {
    return {
        "<string>",
        absl::Substitute("\"$0sandboxed_api/lenval_core.h\"", kIncludePrefix),
    };
  }
  std::string EmitHostPreCall() const override {
    return absl::Substitute(
        "sapi::v::LenVal sapi_tmp_$0($0.data(), $0.size());\n", name_);
  }
  std::string EmitHostPostCall() const override {
    return absl::Substitute(
        "$0.assign(reinterpret_cast<char*>(sapi_tmp_$0.GetData()), "
        "sapi_tmp_$0.GetDataSize());\n",
        name_);
  }
  std::string EmitHostArgs() const override {
    return absl::Substitute("sapi_tmp_$0.PtrBoth()", name_);
  }
  std::string EmitSandboxeeParams() const override {
    return absl::Substitute("sapi::LenValStruct* $0", name_);
  }
  std::string EmitSandboxeeArgs() const override {
    return absl::Substitute("sapi_tmp_$0", name_);
  }
  std::string EmitSandboxeePreCall() const override {
    return absl::Substitute(
        "std::string sapi_tmp_$0(reinterpret_cast<char*>($0->data), "
        "$0->size);\n"
        "free($0->data);\n",
        name_);
  }
  std::string EmitSandboxeePostCall() const override {
    return absl::Substitute(
        "$0->data = strdup(sapi_tmp_$0.c_str());\n"
        "$0->size = sapi_tmp_$0.size();\n",
        name_);
  }
};

// "std::string_view", pretty much the same as "const std::string&".
struct StringViewArg : SandboxedLibraryEmitter::Arg {
  using Arg::Arg;
  std::vector<std::string> Includes() const override {
    return {"<string_view>"};
  }
  std::string EmitHostPreCall() const override {
    return absl::Substitute(
        "sapi::v::Array<char> sapi_tmp_$0(const_cast<char*>($0.data()), "
        "$0.size());\n",
        name_);
  }
  std::string EmitHostArgs() const override {
    return absl::Substitute("sapi_tmp_$0.PtrBefore(), $0.size()", name_);
  }
  std::string EmitSandboxeeParams() const override {
    return absl::Substitute("const char* $0_data, size_t $0_size", name_);
  }
  std::string EmitSandboxeeArgs() const override {
    return absl::Substitute("std::string_view($0_data, $0_size)", name_);
  }
};

absl::StatusOr<std::string> SandboxedLibraryEmitter::EmitSandboxeeHdr(
    const GeneratorOptions& options) const {
  std::string out;
  for (const auto* func : SortedFuncs()) {
    EmitWrapperDecl(out, *func);
    out += ";\n\n";
  }
  return Finalize(out, /*is_header=*/true, /*add_includes=*/true);
}

absl::StatusOr<std::string> SandboxedLibraryEmitter::EmitSandboxeeSrc(
    const GeneratorOptions& options) const {
  std::string out;
  for (const auto* func : SortedFuncs()) {
    EmitFuncDecl(out, *func);
    out += ";\n\n";
    EmitWrapperDecl(out, *func);
    out += " {\n";
    for (const auto& arg : func->args) {
      out += arg->EmitSandboxeePreCall();
    }
    out += "\n";
    if (func->ret) {
      out += "auto sapi_ret_val = ";
    }
    out += absl::Substitute("$0(", func->name);
    for (const auto& arg : func->args) {
      if (&arg != &func->args[0]) {
        out += ", ";
      }
      out += arg->EmitSandboxeeArgs();
    }
    out += ");\n\n";
    for (const auto& arg : func->args) {
      out += arg->EmitSandboxeePostCall();
    }
    if (func->ret) {
      out += func->ret->EmitSandboxeeRet();
    }
    out += "}\n\n";
  }
  return Finalize(out, /*is_header=*/false, /*add_includes=*/true);
}

absl::StatusOr<std::string> SandboxedLibraryEmitter::EmitSandboxeeMain(
    const GeneratorOptions& options) const {
  std::string out;
  for (const auto* func : SortedFuncs()) {
    out += absl::Substitute("extern \"C\" void $0();\n", func->name);
  }
  out += "\nint main() {\n";
  for (const auto* func : SortedFuncs()) {
    out += absl::Substitute("$0();\n", func->name);
  }
  out += "}\n";
  return Finalize(out, /*is_header=*/false, /*add_includes=*/false);
}

absl::StatusOr<std::string> SandboxedLibraryEmitter::EmitHostSrc(
    const GeneratorOptions& options) const {
  std::string out;
  out += absl::Substitute("#include \"$0\"\n", options.out_file);
  out += absl::Substitute(kHostHeader, kIncludePrefix, options.name);
  for (const auto* func : SortedFuncs()) {
    EmitFuncDecl(out, *func);
    out += " {\n";
    out += absl::Substitute("auto* sandbox = $0SandboxImpl::Instance();\n",
                            options.name);
    out += absl::Substitute("$0Api api(sandbox);\n\n", options.name);
    for (const auto& arg : func->args) {
      out += arg->EmitHostPreCall();
    }
    if (func->ret) {
      out += func->ret->EmitRetPreCall();
    }
    out += "\n";
    out += absl::Substitute("sandbox->Check(api.$0$1(", kWrapperPrefix,
                            func->name);
    for (const auto& arg : func->args) {
      if (&arg != &func->args[0]) {
        out += ", ";
      }
      out += arg->EmitHostArgs();
    }
    if (func->ret) {
      if (!func->args.empty()) {
        out += ", ";
      }
      out += func->ret->EmitRetArgs();
    }
    out += "));\n\n";
    for (const auto& arg : func->args) {
      out += arg->EmitHostPostCall();
    }
    if (func->ret) {
      out += func->ret->EmitHostRet();
    }
    out += "}\n\n";
  }
  return Finalize(out, /*is_header=*/false, /*add_includes=*/true);
}

void SandboxedLibraryEmitter::EmitFuncDecl(std::string& out, const Func& func) {
  std::string ret = "void";
  if (func.ret) {
    ret = func.ret->EmitRetType();
  }
  out += absl::Substitute("extern \"C\" $0 $1(", ret, func.name);
  for (const auto& arg : func.args) {
    if (&arg != &func.args[0]) {
      out += ", ";
    }
    out += arg->EmitHostParams();
  }
  out += ")";
}

void SandboxedLibraryEmitter::EmitWrapperDecl(std::string& out,
                                              const Func& func) {
  // Strictly speaking, used/retain attributes are not needed here,
  // but they are required now to work around broken global sapi_library mode
  // (when functions attribute is empty).
  out +=
      absl::Substitute("extern \"C\" __attribute__((used, retain)) void $0$1(",
                       kWrapperPrefix, func.name);
  for (const auto& arg : func.args) {
    if (&arg != &func.args[0]) {
      out += ", ";
    }
    out += arg->EmitSandboxeeParams();
  }
  if (func.ret) {
    if (!func.args.empty()) {
      out += ", ";
    }
    out += func.ret->EmitRetParams();
  }
  out += ")";
}

absl::StatusOr<std::string> SandboxedLibraryEmitter::Finalize(
    const std::string& body, bool is_header, bool add_includes) const {
  std::string out;
  if (is_header) {
    out.append(kHeaderHeader.data(), kHeaderHeader.size());
  }
  out.append(kCommonHeader.data(), kCommonHeader.size());
  if (add_includes) {
    std::vector<std::string> includes(includes_.begin(), includes_.end());
    std::sort(includes.begin(), includes.end());
    for (const auto& inc : includes) {
      out += absl::Substitute("#include $0\n", inc);
    }
  }
  out += "\n";
  out += body;
  out.append(kCommonFooter.data(), kCommonFooter.size());
  return internal::ReformatGoogleStyle("input", out);
}

absl::Status SandboxedLibraryEmitter::AddFunction(clang::FunctionDecl* decl) {
  std::unique_ptr<Arg> ret;
  clang::QualType ret_type = decl->getReturnType();
  if (!ret_type->isVoidType()) {
    ret = Convert("sapi_ret_arg", ret_type);
    if (!ret || ret->EmitRetParams().empty()) {
      return absl::UnimplementedError(absl::Substitute(
          "unsupported return type: $0 ($1)", ret_type.getAsString(),
          ret_type.getCanonicalType().getAsString()));
    }
    for (const std::string& inc : ret->Includes()) {
      includes_.insert(inc);
    }
  }

  std::vector<ArgPtr> args;
  for (size_t i = 0; i < decl->getNumParams(); ++i) {
    const clang::ParmVarDecl* param = decl->getParamDecl(i);
    std::string name = param->getNameAsString();
    if (name.empty()) {
      name = absl::StrFormat("sapi_arg%zu", i);
    }
    clang::QualType type = param->getType();
    std::unique_ptr<Arg> arg = Convert(name, type);
    if (!arg) {
      return absl::UnimplementedError(absl::Substitute(
          "arg $0: unsupported type: $1 ($2)", name, type.getAsString(),
          type.getCanonicalType().getAsString()));
    }
    for (const std::string& inc : arg->Includes()) {
      includes_.insert(inc);
    }
    args.push_back(std::move(arg));
  }

  std::string name =
      clang::ASTNameGenerator(decl->getASTContext()).getName(decl);
  funcs_[name] =
      std::unique_ptr<Func>(new Func{name, std::move(ret), std::move(args)});
  return absl::OkStatus();
}

SandboxedLibraryEmitter::ArgPtr SandboxedLibraryEmitter::Convert(
    absl::string_view name, clang::QualType type) {
  // We are not interested in typedefs.
  type = type.getCanonicalType();
  std::string type_name = type.getAsString();
  if (type->isArithmeticType()) {
    return std::make_unique<ScalarArg>(name, type_name);
  }
  if (type_name == "std::string" ||
      type_name == "class std::basic_string<char>") {
    return std::make_unique<StringArg>(name, type_name);
  }
  if (type_name == "const std::string &" ||
      type_name == "const class std::basic_string<char> &") {
    return std::make_unique<StringConstRefArg>(name, type_name);
  }
  if (type_name == "std::string &" ||
      type_name == "class std::basic_string<char> &") {
    return std::make_unique<StringRefArg>(name, type_name);
  }
  if (type_name == "std::string_view" ||
      type_name == "class std::basic_string_view<char>") {
    return std::make_unique<StringViewArg>(name, type_name);
  }
  return {};
}

std::vector<const SandboxedLibraryEmitter::Func*>
SandboxedLibraryEmitter::SortedFuncs() const {
  std::vector<const Func*> sorted;
  sorted.reserve(funcs_.size());
  for (const auto& func : funcs_) {
    sorted.push_back(func.second.get());
  }
  std::sort(sorted.begin(), sorted.end(),
            [](const Func* a, const Func* b) { return a->name < b->name; });
  sorted.erase(std::unique(sorted.begin(), sorted.end(),
                           [](const Func* a, const Func* b) {
                             return a->name == b->name;
                           }),
               sorted.end());
  return sorted;
}

SandboxedLibraryEmitter::~SandboxedLibraryEmitter() = default;
SandboxedLibraryEmitter::Func::~Func() = default;

}  // namespace sapi
